use crate::{AppError, AppState};
use axum::{
    async_trait,
    body::Body,
    http::{header, HeaderValue, Request},
    response::{IntoResponse, Response},
};
use base64::{prelude::BASE64_STANDARD, Engine};
use chrono::Days;
use common::utils::current_datetime;
use cookie::Cookie;
use hmac::{Hmac, Mac};
use loginmanager::{DecodeRequest, LoginInfo};
use sea_orm::{prelude::Expr, ActiveModelTrait, ColumnTrait, EntityTrait, QueryFilter, Set};
use sha2::Sha256;

type HmacSha256 = Hmac<Sha256>;

fn hmacsha256_base64(key: &[u8], content: &[u8]) -> String {
    let mut hasher = HmacSha256::new_from_slice(key).expect("");
    hasher.update(content);
    BASE64_STANDARD.encode(hasher.finalize().into_bytes())
}

#[derive(Clone)]
pub struct AneBackend {
    state: AppState,
    session_key: String,
    secret: String,
}

impl AneBackend {
    pub fn new(state: AppState) -> Self {
        Self {
            state,
            session_key: "__session".to_owned(),
            secret: "secret".to_owned(),
        }
    }

    #[inline]
    fn generate_key(id: &str) -> String {
        format!("__session_{id}")
    }
}

#[async_trait]
impl DecodeRequest<Request<Body>, Response> for AneBackend {
    async fn decode(&self, req: &mut Request<Body>) -> Result<Option<String>, Response> {
        /// 依次从header和cookie中读取认字符串
        let mut access_token = req
            .headers()
            .get(header::AUTHORIZATION)
            .and_then(|s| s.to_str().ok())
            .and_then(|s| {
                if s.starts_with("Bearer") {
                    Some(s[6..].trim().to_owned())
                } else {
                    None
                }
            });
        if access_token.is_none() {
            let mut cookie_find = "".to_owned();
            let cookie_name = format!("{}=", self.session_key);
            for hdr in req.headers().get_all(header::COOKIE) {
                if let Ok(s) = hdr.to_str() {
                    for cookie_str in s.split(';').map(|s| s.trim()) {
                        if cookie_str.starts_with(&cookie_name) {
                            cookie_find = cookie_str[self.session_key.len() + 1..].to_string();
                        }
                    }
                }
            }
            if !cookie_find.is_empty() {
                access_token = Some(cookie_find);
            }
        };

        if let Some(token) = access_token {
            /// 认证字符全格式 token.base64(hmac(token))
            let tokens = token.split(".").collect::<Vec<_>>();
            if tokens.len() != 2 {
                return Ok(None);
            }
            let token = tokens[0];
            if hmacsha256_base64(self.secret.as_bytes(), token.as_bytes()) != tokens[1] {
                return Ok(None);
            }
            /// 将token保存，退出时需要使用
            let login_info = req.extensions().get::<LoginInfo>().unwrap();
            login_info.set_ext(Some(token.to_owned()));
            /// 从cahce中读取token对应的user_id, 并增加数据库中字段过期的时间
            let cache_key = Self::generate_key(token);
            if let Some(v) = self.state.cache().get(&cache_key).await {
                crate::model::entity::sys_auth::Entity::update_many()
                    .col_expr(
                        crate::model::entity::sys_auth::Column::ExpiresAt,
                        Expr::value(current_datetime().checked_add_days(Days::new(1)).unwrap()),
                    )
                    .col_expr(
                        crate::model::entity::sys_auth::Column::UpdatedAt,
                        Expr::value(current_datetime()),
                    )
                    .filter(crate::model::entity::sys_auth::Column::AccessToken.eq(token))
                    .exec(self.state.db())
                    .await;
                return Ok(Some(v));
            }
            /// 从数据库中读取token对应的user_id, 并将结果保存到cache中, 有效期 30min
            if let Some(s) = crate::model::entity::sys_auth::Entity::find()
                .filter(crate::model::entity::sys_auth::Column::AccessToken.eq(token))
                .filter(crate::model::entity::sys_auth::Column::Deleted.eq(false))
                .filter(crate::model::entity::sys_auth::Column::ExpiresAt.gt(current_datetime()))
                .one(self.state.db())
                .await
                .map_err(|e| Into::<AppError>::into(e).into_response())?
            {
                self.state
                    .cache()
                    .set_ex(&cache_key, s.user_id.clone(), 1000 * 1800)
                    .await;
                return Ok(Some(s.user_id));
            }
        }
        return Ok(None);
    }

    async fn update(&self, res: &mut Response) {
        if let Some(login_info) = res.extensions().get::<LoginInfo>() {
            if let Some(key) = login_info.login_key() {
                let s = crate::model::entity::sys_auth::ActiveModel {
                    access_token: Set(gen_table::util::ulid_str()),
                    user_id: Set(key),
                    deleted: Set(false),
                    expires_at: Set(current_datetime().checked_add_days(Days::new(1)).unwrap()),
                    created_at: Set(current_datetime()),
                    updated_at: Set(current_datetime()),
                    ..Default::default()
                };
                let s = s.insert(self.state.db()).await.unwrap();
                let hash = hmacsha256_base64(self.secret.as_bytes(), s.access_token.as_bytes());
                let mut cookie =
                    Cookie::new(&self.session_key, format!("{}.{}", s.access_token, hash));
                cookie.set_path("/");
                res.headers_mut().append(
                    header::SET_COOKIE,
                    HeaderValue::from_str(&cookie.to_string()).unwrap(),
                );
            } else if login_info.is_logout() {
                if let Some(token) = login_info.ext() {
                    crate::model::entity::sys_auth::Entity::update_many()
                        .col_expr(
                            crate::model::entity::sys_auth::Column::Deleted,
                            Expr::value(true),
                        )
                        .col_expr(
                            crate::model::entity::sys_auth::Column::UpdatedAt,
                            Expr::value(current_datetime()),
                        )
                        .filter(crate::model::entity::sys_auth::Column::AccessToken.eq(&token))
                        .exec(self.state.db())
                        .await;
                    self.state.cache().del(&Self::generate_key(&token)).await;
                }
            };
        }
    }
}
